import os
from CalTorque import CalTorque
import numpy
import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot as PLT
from matplotlib import dates
from datetime import time


def CompareBeforeAfterBreak(data_files, bins = 400, max_d = 400, outdir = 'AveragePlots/'):
    telescope = {}
    telescope['before'] = []
    telescope['after'] = []
    
    for file in data_files:
        print file
        data = CalTorque(file)
        if data['source'] != 'weakTh-228': 
            continue
        if data['station'] in ['S8', 'S11', 'Other']:
            continue
        date = int(''.join(str(data['datetime'][0]).split(' ')[0].split('-')))
        print date
        if date <= 20130222:
            telescope['before'].append([data['pos_offsetfix'], data['torque_inoz']])
            # print 'before'
        if date > 20130222:
            telescope['after'].append([data['pos_offsetfix'], data['torque_inoz']])
            # print 'after'

    for t in telescope:
        x_inputs = numpy.zeros(bins)
        X = numpy.linspace(0,max_d,bins)
        Y = numpy.zeros(bins)
        for i, k in enumerate(telescope[t]):
            # k[0] = x, k[1] = y
            for j, x in enumerate(X):
                pos_torque = []
                x0 = x-(max_d/bins)/2.
                x1 = x+(max_d/bins)/2.
                for i, p in enumerate(k[0]):
                    if p > x0 and p <= x1:
                        pos_torque.append(k[1][i])
                if len(pos_torque)>0:
                    if max(pos_torque) == 0.: continue
                    Y[j] += max(pos_torque)
                    x_inputs[j]+=1.
    
        for l, n in enumerate(x_inputs):
            if n == 0: continue
            Y[l] = Y[l]/n
        telescope[t] = [X, Y]
        
    
    fig = PLT.figure(figsize=(15,8),dpi=150)
    box = [0.14, 0.14, 0.76, 0.76]
    ax = fig.add_axes(box)
    ax.set_xlabel('Position (cm)')
    ax.set_ylabel('Torque (in.oz)')
    
    miny = 1
    maxx = 400
    
    ## add box at telescope location
    ax.axvspan(20.8, 33.3, color = 'magenta', alpha = 0.2)
    ax.annotate('Telescope Travel', xy = (20.8, miny + 2), color = 'magenta')
            
    ## Add station locations
    if maxx >= 328.9:
        ax.axvline(338.9, color = 'r', linewidth = 1, linestyle = '-')
        ax.annotate('S2', xy = (340.9, miny * 0.8), color = 'r')
    if maxx >= 382.6:
        ax.axvline(392.6, color = 'r', linewidth = 1, linestyle = '-')
        ax.annotate('S5', xy = (394.6, miny * 0.8), color = 'r')
    if maxx >= 435.7:
        ax.axvline(445.7, color = 'r', linewidth = 1, linestyle = '-')
        ax.annotate('S8', xy = (447.7, miny * 0.8), color = 'r')
    if maxx >= 491.7:
        ax.axvline(501.7, color = 'r', linewidth = 1, linestyle = '-')
        ax.annotate('S11', xy = (503.7, miny * 0.8), color = 'r')
    if maxx >= 595.8:
        ax.axvline(605.8, color = 'r', linewidth = 1, linestyle = '-')
        ax.annotate('S17', xy = (607.8, miny * 0.8), color = 'r')
    
    ## Add bend locations
    if maxx >= 47.8:
        ax.axvspan(47.8, 93.5, color = 'g', alpha = 0.1)
        ax.annotate('Bellows', xy = (48, miny), color = 'g')
    if maxx >= 136.4:
        ax.axvspan(136.4, 144.6, color = 'g', alpha = 0.1)
        ax.annotate('First outer bend', xy = (136, miny), color = 'g')
    if maxx >= 196.5:
        ax.axvspan(196.5, 218.0, color = 'g', alpha = 0.1)
        ax.annotate('The "S"', xy = (212, miny), color = 'g')
    if maxx >= 228.3:
        ax.axvspan(228.3, 241.5, color = 'g', alpha = 0.1)
    if maxx >= 274.6:
        ax.axvspan(274.6, 284.6, color = 'g', alpha = 0.1)
        ax.annotate('Bend 1', xy = (274, miny), color = 'g')
    if maxx >= 309.0:
        ax.axvspan(309.0, 318.9, color = 'g', alpha = 0.1)
        ax.annotate('Bend 2', xy = (309, miny), color = 'g')
    if maxx >= 358.8:
        ax.axvspan(358.8, 370.8, color = 'g', alpha = 0.1)
        ax.annotate('Bend 3', xy = (359, miny), color = 'g')
    if maxx >= 414.5:
        ax.axvspan(414.5, 426.5, color = 'g', alpha = 0.1)
        ax.annotate('Bend 4', xy = (414, miny), color = 'g')
    if maxx >= 445.3:
        ax.axvspan(445.3, 455.3, color = 'g', alpha = 0.1)
        ax.annotate('Bend 5', xy = (445, miny), color = 'g')
    if maxx >= 467.5:
        ax.axvspan(467.5, 479.5, color = 'g', alpha = 0.1)
        ax.annotate('Bend 6', xy = (467, miny), color = 'g')
    if maxx >= 523.8:
        ax.axvspan(523.8, 535.8, color = 'g', alpha = 0.1)
        ax.annotate('Bend 7', xy = (523, miny), color = 'g')
    if maxx >= 571.6:
        ax.axvspan(571.6, 583.6, color = 'g', alpha = 0.1)
        ax.annotate('Bend 8', xy = (571, miny), color = 'g')
    if maxx >= 627.9:
        ax.axvspan(627.9, 639.9, color = 'g', alpha = 0.1)
        ax.annotate('Bend 9', xy = (627, miny), color = 'g')
    if maxx >= 652.7:
        ax.axvspan(652.7, 660.7, color = 'g', alpha = 0.1)
        ax.annotate('Bend 10', xy = (652, miny), color = 'g')
    if maxx >= 678.5:
        ax.axvspan(678.5, 681.1, color = 'g', alpha = 0.1)
        ax.annotate('Bend 11', xy = (678, miny), color = 'g')
    
    ax.plot(telescope['before'][0], telescope['before'][1], '-', linewidth = 2, c = 'b', label = 'Average Torque Before Break')
    ax.plot(telescope['after'][0], telescope['after'][1], '-', linewidth = 2, c = 'r', label = 'Average Torque After Break')
    h, l = ax.get_legend_handles_labels()
    ax.grid()
    PLT.title('Average Torque Before and After Telescope Break')
    PLT.legend(h, l, 'upper left')
    PLT.savefig(outdir + 'Before_After_Telescope_Break.png')